import HomologFinderUtility, FileUtility
from datamodel import HomologFinderConfig
from Bio import SeqIO
import re
import logging


class STRINGUtility:
    STRING_HOME = HomologFinderConfig.STRING_HOME
    MULTI_COG = STRING_HOME + "multi.cog.proteins"
    COG_MAPPINGS = STRING_HOME + "COG.mappings.v8.2.txt"
    STRING_SEQS = STRING_HOME + "protein.sequences.v8.2.fa"
    STRING_DB = STRING_HOME + "STRING.v.8.2"
    STRING_LINKS = STRING_HOME + "COG.links.detailed.v8.2.txt"

    stringDBIndex = None

    def __init__(self):
        logging.debug("initializing STRINGUtility")
        #self.__class__.stringDBIndex = SeqIO.index(self.STRING_SEQS, 'fasta')
        logging.debug("STRING DB Index created")

        
    def getMultiCOGSequences(self, sequences):
        """Checks the multi.cog.proteins file to see if a sequence has more than one COG"""
        logging.info("Checking the list of sequences for ones with multiple COGs assigned")
        multiList = list()

        try:
            f = open(self.MULTI_COG, 'r', 1)
            str = "|".join(sequences)
            str += '\t'
            pat = re.compile(str)
            for line in f:

                match = re.match(line)
                if match:
                    multiList.append(match.group(0))
                if(len(multiList) == len(sequences)): #all have been found, may as well break
                    break



        finally:
            FileUtility.safeFileClose(f)
            return multiList

    def getSequencesMatchingCOG(self, cog):
        """Parses all the hits for a particular cog from the STRING mappings file"""
        matches = list() #FileUtility.parseFile(self.__COG_MAPPINGS, r'^(\S+)\t\S+\t\S+\t' + cog)
        try:
           f = open(self.COG_MAPPINGS, 'r', 1)
           for line in f:
               if cog in line:
                   match = re.match('\S+', line)
                   if match:
                       matches.append(match.group(0))
        finally:
            FileUtility.safeFileClose(f)
        return matches



    def getFastaList(self, seqIds, fastaFile):
        """Gets a list of fastas from a list of sequence ids"""
        fastaList = list()
        self.stringDBIndex = SeqIO.index(fastaFile, 'fasta')

        for seq in seqIds:
            fastaList.append(self.stringDBIndex[seq])

        return fastaList

def filterSequencesByCOG(singleCOG, cog):
        """Filters a list to those which match a specific COG"""
        logging.info("Identifying sequences assigned to %s" % cog)
        matches = list()
        try:
            f = open(STRINGUtility.COG_MAPPINGS, 'r', 1)
            str = '('
            str += '|'.join(singleCOG)
            str += ')'
            pat = re.compile(r'%s\t\S+\t\S+\t%s' % (str, cog))
            for line in f:
                match = pat.match(line)
                if match:
                    matches.append(match.group(1))
                if(len(matches) == len(singleCOG)):
                    break
        finally:
            FileUtility.safeFileClose(f)

        return matches

def checkForGeneFusion(sequences, cog):
    """checks to see if sequences for a particular cog are fused, if they don't match the specified
        cog
    """
    raise Exception("Not yet implemented!")
    matches = dict()
    try:
        f = open(STRINGUtility.MULTI_COG, 'r', 1)
        str = '('
        str += '|'.join(sequences)
        str += ')\t(\S*' + cog + '\S*)'
        pat = re.compile(str)
        for line in f:
           match = pat.match(line)
           if match:
               tmp = re.split(',', match.group(2))
               #Don't want to add them if they KOGs or NOGs
               cogs = []
               for cog in tmp:
                   if "KOG" not in cog and "NOG" not in cog:
                       cogs.append(cog)
               matches[match.group(1)] = cogs
    finally:
        FileUtility.safeFileClose(f)

    #Now we have list of multiple matches, we can check to see if those cogs are ever fused
    try:
        f = open(STRINGUtility.STRING_LINKS, 'r', 1)
        for key in matches.keys():
            cogs = matches[key]
            #if cogs contains more than 2 groups, we need to split it into pairs of cog in question
            #plus one other.
            #match can either be cog1\tcog2 or cog2\tcog1
            # (?:...) makes a non-matching group
            pass
    finally:
        FileUtility.safeFileClose(f)

    pass

 